#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""计算矢量点在平面上的欧氏距离

通过给定缓冲区和像元大小, 以栅格的形式展示矢量点在给定范围内的欧氏距离。
其中 *像元大小* 为必填项; 缓冲距离默认为 1000 米。

Example:
    $ python euclidean-distance-vector.py -distunits 81.39 -bufferdist 2000
"""

import sys
import math
import numpy as np
from osgeo import gdal, ogr, osr

def compute_distance(point_a, point_b):
    a_x, a_y = point_a
    b_x, b_y = point_b
    distance = math.sqrt(math.pow(b_x - a_x, 2) + math.pow(b_y - a_y, 2))

    return distance

def find_minial_distance(target_point, coordinates):
    p_x, p_y = target_point

    minimal = None
    for point in coordinates:
        distance = compute_distance(target_point, point)

        if minimal is None:
            minimal = distance
        elif distance < minimal:
            minimal = distance

    return minimal

def offset2coordinate(geotransform, offset):
    originX = geotransform[0]
    originY = geotransform[3]
    pixelWidth = geotransform[1]
    pixelHeight = geotransform[5]

    x_offset, y_offset = offset

    x = x_offset * pixelWidth + originX
    y = y_offset * pixelHeight + originY

    return (x, y)

def array2raster_shp(shape_fn, new_raster_fn, geotransform, array):
    shapefile = ogr.Open(shape_fn)
    layer = shapefile.GetLayer()

    rows, cols = array.shape

    driver = gdal.GetDriverByName('GTiff')
    out_raster = driver.Create(new_raster_fn, cols, rows, 1, gdal.GDT_Float32)
    out_raster.SetGeoTransform(geotransform)
    out_band = out_raster.GetRasterBand(1)
    out_band.WriteArray(array)

    srs = layer.GetSpatialRef()
    out_raster.SetProjection(srs.ExportToWkt())
    out_band.FlushCache()

def main(shape_fn, new_raster_fn, buffer_dist, dist_units):
    datasource = ogr.Open(shape_fn)
    layer = datasource.GetLayer()

    point_x = []
    point_y = []
    points = []
    for feature in layer:
        geom = feature.GetGeometryRef()
        x, y, _ = geom.GetPoint()

        point_x.append(x)
        point_y.append(y)

        points.append((x, y))

    x_min = min(point_x) - buffer_dist
    x_max = max(point_x) + buffer_dist
    y_min = min(point_y) - buffer_dist
    y_max = max(point_y) + buffer_dist

    rows = int((x_max - x_min) / dist_units) + 1
    cols = int((y_max - y_min) / dist_units) + 1

    geotransform = (x_min, dist_units, 0, y_max, 0, -dist_units)

    distance_array = np.zeros((rows, cols))

    for x, y in np.ndindex(distance_array.shape):
        geo_point = offset2coordinate((x_min, dist_units, 0, y_max, 0, -dist_units), (x, y))
        distance_array[x, y] = find_minial_distance(geo_point, points)

    for x, y in points:
        x_offset = int((x - x_min) / dist_units)
        y_offset = int((y - y_max) / -dist_units)

        distance_array[x_offset, y_offset] = 0

    array2raster_shp(shape_fn, new_raster_fn, geotransform, distance_array.T)

if __name__ == "__main__":
    shape_fn = None
    new_raster_fn = None
    buffer_dist = None
    dist_units = None

    argv = sys.argv

    i = 1
    while i < len(argv):
        arg = argv[i]

        if arg == '-distunits':
            i = i + 1
            dist_units = float(argv[i])

        elif arg == '-bufferdist':
            i = i + 1
            buffer_dist = float(argv[i])

        elif shape_fn is None:
            shape_fn = argv[i]

        elif new_raster_fn is None:
            new_raster_fn = argv[i]

        else:
            print("python script.py [-distunits n] [-bufferdist n] shape.shp out.tif")
            sys.exit(1)

        i = i + 1

    if buffer_dist is None:
        buffer_dist = 1000

    if new_raster_fn is None:
        new_raster_fn = 'OUT.tif'

    if dist_units is None:
        print("需要提供像元大小!")
        sys.exit(1)

    main(shape_fn, new_raster_fn, buffer_dist, dist_units)

